#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"

using namespace AscendC;

constexpr uint32_t BUFFER_NUM_T_ONE = 1;
constexpr uint32_t BUFFER_NUM = 1;
constexpr int32_t nnLength = 4096; // 64*64
constexpr int32_t nLength = 64;
constexpr int32_t nSize = 64;
class KernelWKV7TOne
{
public:
    __aicore__ inline KernelWKV7TOne() {}
    __aicore__ inline void Init(uint32_t bSize, uint32_t tSize, uint32_t hSize,
                                GM_ADDR k, GM_ADDR v, GM_ADDR w, GM_ADDR r, GM_ADDR a, GM_ADDR b, 
                                GM_ADDR h0, GM_ADDR o,  GM_ADDR ht, TPipe* pipeIn)
    {
        this->tSize = tSize;
        this->bSize = bSize;
        this->hSize = hSize;
        this->hnLength = hSize * nLength;
        this->hnnLength = hSize * nnLength;
        this->thnLength = 2560;//hSize * nSize;
        this->bhLength = bSize * hSize;
        this->pipe = pipeIn;
        this->tileLength = 2;
        int32_t blockNum = GetBlockNum();
        int32_t coreId = GetBlockIdx();
        int32_t loopCount, start_index, end_index;
        #if __CCE_AICORE__ == 200
            blockNum = 15;
            uint32_t unitCount = (bhLength + 3) / 4; // 划分成4份
            if(coreId >= 10){
                // 后1份 分到 后5个核
                loopCount = (unitCount + 4) / 5;
                start_index = 3 * unitCount + (coreId - 10) * loopCount;
                end_index = start_index + loopCount;
                if(end_index > bhLength){
                    end_index = bhLength;
                }
            }else{
                //前3份 分到 前10个核
                unitCount *= 3;
                loopCount =  (unitCount + 9) / 10;
                start_index = coreId * loopCount;
                end_index = start_index + loopCount;
                tileLength = 3;
                if(end_index > unitCount || coreId == 9){
                    end_index = unitCount;
                }
            }
        #else
            loopCount = bhLength / blockNum;
            start_index = loopCount * coreId;
            end_index = start_index + loopCount;
            if(end_index > bhLength || coreId == blockNum - 1){
                end_index = bhLength;
            }
        #endif
        if(start_index > end_index){
            start_index = end_index;
        }
        indexLength = end_index - start_index;
        if(indexLength < tileLength){
            tileLength = indexLength;
        }
        int32_t offsetBig = start_index * nnLength;
        int32_t offsetSmall = start_index * nLength;
        kGm.SetGlobalBuffer((__gm__ half *)k + offsetSmall, indexLength * nLength);
        vGm.SetGlobalBuffer((__gm__ half *)v + offsetSmall, indexLength * nLength);
        wGm.SetGlobalBuffer((__gm__ half *)w + offsetSmall, indexLength * nLength);
        rGm.SetGlobalBuffer((__gm__ half *)r + offsetSmall, indexLength * nLength);
        oGm.SetGlobalBuffer((__gm__ half *)o + offsetSmall, indexLength * nLength);
        aGm.SetGlobalBuffer((__gm__ half *)a + offsetSmall, indexLength * nLength);
        bGm.SetGlobalBuffer((__gm__ half *)b + offsetSmall, indexLength * nLength);
        h0Gm.SetGlobalBuffer((__gm__ half *)h0 + offsetBig, indexLength * nnLength);
        htGm.SetGlobalBuffer((__gm__ half *)ht + offsetBig, indexLength * nnLength);
        pipe->InitBuffer(inQueueKVWRAB, BUFFER_NUM_T_ONE, 6 * tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueH, BUFFER_NUM_T_ONE, tileLength * nnLength * sizeof(half));
        pipe->InitBuffer(outQueueH, BUFFER_NUM_T_ONE, tileLength * nnLength * sizeof(half));
        pipe->InitBuffer(outQueueO, BUFFER_NUM_T_ONE, tileLength * nLength * sizeof(half));

        pipe->InitBuffer(vkBuf, nnLength * sizeof(half));
        pipe->InitBuffer(habBuf, nnLength * sizeof(half));
        pipe->InitBuffer(tmpBuf, nnLength * sizeof(float));
        pipe->InitBuffer(calcBuf, nnLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        
        for(int32_t index = 0; index < indexLength; index+=tileLength){
            if(index + tileLength > indexLength){
                tileLength = indexLength - index;
            }
            CopyInTOne(index * nnLength, index * nLength);
            ComputeTOne();
            CopyOutTOne(index * nnLength, index * nLength);
        }
    }

private:
    __aicore__ inline void CopyInTOne(int32_t offset1, int32_t offset)
    {
        // B, H, N, N(40, 64, 64)
        // b, h
        auto h0Local = inQueueH.AllocTensor<half>();
        auto kLocal = inQueueKVWRAB.AllocTensor<half>();
        auto vLocal = kLocal[tileLength * nLength];
        auto wLocal = vLocal[tileLength * nLength];
        auto rLocal = wLocal[tileLength * nLength];
        auto aLocal = rLocal[tileLength * nLength];
        auto bLocal = aLocal[tileLength * nLength];

        DataCopy(h0Local, h0Gm[offset1], {tileLength, 256, 0, 0}); // 64*64*2 / 32
        DataCopy(wLocal, wGm[offset], {tileLength, 4, 0, 0});
        DataCopy(vLocal, vGm[offset], {tileLength, 4, 0, 0});
        DataCopy(kLocal, kGm[offset], {tileLength, 4, 0, 0});
        DataCopy(aLocal, aGm[offset], {tileLength, 4, 0, 0});
        DataCopy(bLocal, bGm[offset], {tileLength, 4, 0, 0});
        DataCopy(rLocal, rGm[offset], {tileLength, 4, 0, 0});
        
        inQueueH.EnQue(h0Local);
        inQueueKVWRAB.EnQue(kLocal);
    }

    __aicore__ inline void ComputeTOne(){
        LocalTensor<half> orih0Local = inQueueH.DeQue<half>();
        LocalTensor<half> orikLocal = inQueueKVWRAB.DeQue<half>();
        LocalTensor<half> orivLocal = orikLocal[tileLength * nLength];
        LocalTensor<half> oriwLocal = orivLocal[tileLength * nLength];
        LocalTensor<half> orirLocal = oriwLocal[tileLength * nLength];
        LocalTensor<half> oriaLocal = orirLocal[tileLength * nLength];
        LocalTensor<half> oribLocal = oriaLocal[tileLength * nLength];

        LocalTensor<half> vkLocal = vkBuf.Get<half>();
        LocalTensor<half> habLocal = habBuf.Get<half>();
        LocalTensor<float> tmpFp = tmpBuf.Get<float>();
        LocalTensor<float> calcFp = calcBuf.Get<float>();

        LocalTensor<half> orioLocal = outQueueO.AllocTensor<half>();
        LocalTensor<half> orihtLocal = outQueueH.AllocTensor<half>();

        LocalTensor<half> h0Local = orih0Local;
        LocalTensor<half> wLocal = oriwLocal;
        LocalTensor<half> rLocal = orirLocal;
        LocalTensor<half> aLocal = oriaLocal;
        LocalTensor<half> bLocal = oribLocal;
        LocalTensor<half> vLocal = orivLocal;
        LocalTensor<half> kLocal = orikLocal;

        LocalTensor<half> oLocal = orioLocal;
        LocalTensor<half> htLocal = orihtLocal;

        Exp(oriwLocal, oriwLocal, this->tileLength * nLength);
        #if __CCE_AICORE__ == 200
            uint64_t dstLocalList[4][16];
            for(int j = 0; j < 4; j++){
                for (int i = 0; i < 16; i++) {
                    dstLocalList[j][i] = (uint64_t)(calcFp[i * 32 + j * 8].GetPhyAddr());
                }
            }
            uint64_t srcLocalList[16];
            for (int i = 0; i < 16; i++) {
                srcLocalList[i] = (uint64_t)(tmpFp[0].GetPhyAddr());
            }
        #endif
        for(int i = 0; i < tileLength; i++){
            Mul(habLocal, h0Local, aLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            Cast(tmpFp, habLocal, RoundMode::CAST_NONE, nnLength);
            WholeReduceSum(tmpFp, tmpFp, nLength, nSize, 1, 1, 8); // get sa

            #if __CCE_AICORE__ == 200
                TransDataTo5HD<float>(dstLocalList[0], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[1], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[2], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[3], srcLocalList, {false, false, 8, 64, 1});
                Cast(htLocal, calcFp, RoundMode::CAST_NONE, nnLength);
                Mul(habLocal, htLocal, bLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });

                Cast(tmpFp, vLocal, RoundMode::CAST_NONE, nLength);
                TransDataTo5HD<float>(dstLocalList[0], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[1], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[2], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[3], srcLocalList, {false, false, 8, 64, 1});
                Cast(htLocal, calcFp, RoundMode::CAST_NONE, nnLength);
            #else
                Brcb(calcFp, tmpFp, 8, {1, 8});
                Brcb(tmpFp, calcFp, 64, {1, 8});
                Cast(htLocal, tmpFp, RoundMode::CAST_NONE, nnLength);
                Mul(habLocal, htLocal, bLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });

                Cast(tmpFp, vLocal, RoundMode::CAST_NONE, nLength);
                Brcb(calcFp, tmpFp, 8, {1, 8});
                Brcb(tmpFp, calcFp, 64, {1, 8});
                Cast(htLocal, tmpFp, RoundMode::CAST_NONE, nnLength);
            #endif

            Mul(vkLocal, htLocal, kLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            
            Mul(htLocal, h0Local, wLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 }); // h0 = h0 * w
            Add(htLocal, htLocal, vkLocal, nnLength); // h0 = h0 + v*k
            Add(htLocal, htLocal, habLocal, nnLength); // h0 = h0 + hab

            // n,n n,1 -> n,1
            Mul(h0Local, htLocal, rLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            Cast(tmpFp, h0Local, RoundMode::CAST_NONE, nnLength);
            WholeReduceSum(tmpFp, tmpFp, nLength, nSize, 1, 1, 8);
            Cast(oLocal, tmpFp, RoundMode::CAST_NONE, nLength);
            h0Local = h0Local[nnLength];
            wLocal = wLocal[nLength];
            rLocal = rLocal[nLength];
            aLocal = aLocal[nLength];
            bLocal = bLocal[nLength];
            vLocal = vLocal[nLength];
            kLocal = kLocal[nLength];
            oLocal = oLocal[nLength];
            htLocal = htLocal[nnLength];
        }
        
        inQueueKVWRAB.FreeTensor(orikLocal);
        inQueueH.FreeTensor(orih0Local);

        outQueueH.EnQue<half>(orihtLocal);
        outQueueO.EnQue<half>(orioLocal);
    }

    __aicore__ inline void CopyOutTOne(int32_t offset1, int32_t offset)
    {
        auto htLocal = outQueueH.DeQue<half>();
        auto oLocal = outQueueO.DeQue<half>();
        DataCopy(htGm[offset1], htLocal, {tileLength, 256, 0, 0});
        DataCopy(oGm[offset], oLocal, {tileLength, 4, 0, 0});
        outQueueH.FreeTensor(htLocal);
        outQueueO.FreeTensor(oLocal);
    }
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM_T_ONE> inQueueKVWRAB, inQueueH;
    TQue<QuePosition::VECOUT, BUFFER_NUM_T_ONE> outQueueO, outQueueH;
    GlobalTensor<half> kGm, vGm, wGm, rGm, aGm, bGm, oGm, h0Gm, htGm;
    TBuf<QuePosition::VECCALC> vkBuf, habBuf, tmpBuf, calcBuf;
    uint32_t tSize, bSize, thnLength, bhLength, indexLength, hSize, hnLength, hnnLength;
    uint16_t tileLength;
};

class KernelWKV7
{
public:
    __aicore__ inline KernelWKV7() {}
    __aicore__ inline void Init(uint32_t bSize, uint32_t tSize, uint32_t hSize,
                                GM_ADDR k, GM_ADDR v, GM_ADDR w, GM_ADDR r, GM_ADDR a, GM_ADDR b, 
                                GM_ADDR h0, GM_ADDR o,  GM_ADDR ht, TPipe* pipeIn)
    {
        this->tSize = tSize;
        this->bSize = bSize;
        this->thnLength = tSize * 2560;//hSize * nSize;
        this->hSize = hSize;
        this->hnLength = hSize * nLength;
        this->hnnLength = hSize * nnLength;
        this->bhLength = bSize * hSize;
        this->tileLength = 2;
        #if __CCE_AICORE__ == 200
            if(bSize == 1){
                this->tileLength = 3;
            }else{
                this->tileLength = 1;
            }
        #endif
        this->pipe = pipeIn;
        kGm.SetGlobalBuffer((__gm__ half *)k);
        vGm.SetGlobalBuffer((__gm__ half *)v);
        wGm.SetGlobalBuffer((__gm__ half *)w);
        rGm.SetGlobalBuffer((__gm__ half *)r);
        oGm.SetGlobalBuffer((__gm__ half *)o);
        aGm.SetGlobalBuffer((__gm__ half *)a);
        bGm.SetGlobalBuffer((__gm__ half *)b);
        h0Gm.SetGlobalBuffer((__gm__ half *)h0);
        htGm.SetGlobalBuffer((__gm__ half *)ht);
        pipe->InitBuffer(inQueueW, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueR, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueK, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueV, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueA, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueB, BUFFER_NUM, this->tileLength * nLength * sizeof(half));
        pipe->InitBuffer(inQueueH, BUFFER_NUM, this->tileLength * nnLength * sizeof(half));
        pipe->InitBuffer(outQueueH, BUFFER_NUM, this->tileLength * nnLength * sizeof(half));
        pipe->InitBuffer(outQueueO, BUFFER_NUM, this->tileLength * nLength * sizeof(half));

        pipe->InitBuffer(vkBuf, nnLength * sizeof(half));
        pipe->InitBuffer(habBuf, nnLength * sizeof(half));
        pipe->InitBuffer(tmpBuf, nnLength * sizeof(float));
        pipe->InitBuffer(calcBuf, nnLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        int32_t blockNum = GetBlockNum();
        #if __CCE_AICORE__ == 200
            blockNum = 15;
        #endif
        int32_t coreId = GetBlockIdx();
        int32_t loopCount = bhLength / blockNum;
        int32_t lastCount = bhLength - loopCount * blockNum;
        int32_t start_index = loopCount * coreId;
        if(coreId < lastCount){
            start_index += coreId;
            loopCount += 1;
        }else{
            start_index += lastCount;
        }
        int32_t end_index = start_index + loopCount;
        if(end_index > bhLength || coreId == blockNum - 1){
            end_index = bhLength;
        }
        int32_t bIndex = start_index / hSize;
        int32_t hIndex = start_index % hSize;
        int32_t hOffset2 = hIndex * nLength;
        int32_t bOffset2 = bIndex * thnLength;

        for(int32_t index = start_index; index < end_index; index+=tileLength, hIndex+=tileLength, hOffset2 += tileLength * nLength){
            if(hIndex >= hSize){
                hIndex = 0;
                bIndex++;
                hOffset2 = 0;
                bOffset2 += thnLength;
            }
            int32_t offset1 = index * nnLength;
            int32_t offset2 = bOffset2 + hOffset2;
            if(index + tileLength > end_index){
                tileLength = end_index - index;
            }
            CopyInH0(offset1);
            for(int tIndex = 0; tIndex < tSize; tIndex++, offset2 += hnLength){
                CopyIn(offset2);
                Compute();
                CopyOut(offset2);
            }
            CopyH0ToHt();
            CopyOutHt(offset1);
        }
    }

private:

    __aicore__ inline void CopyInH0(int32_t offset)
    {
        // B, H, N, N(40, 64, 64)
        // b, h
        auto h0Local = inQueueH.AllocTensor<half>();
        DataCopy(h0Local, h0Gm[offset], this->tileLength * nnLength);
        inQueueH.EnQue(h0Local);
    }

    __aicore__ inline void CopyIn(int32_t offset)
    {
        // B T H N
        // b t h
        auto wLocal = inQueueW.AllocTensor<half>();
        DataCopy(wLocal, wGm[offset], this->tileLength * nLength);
        inQueueW.EnQue(wLocal);

        auto vLocal = inQueueV.AllocTensor<half>();
        DataCopy(vLocal, vGm[offset], this->tileLength * nLength);
        inQueueV.EnQue(vLocal);

        auto kLocal = inQueueK.AllocTensor<half>();
        DataCopy(kLocal, kGm[offset], this->tileLength * nLength);
        inQueueK.EnQue(kLocal);

        auto aLocal = inQueueA.AllocTensor<half>();
        DataCopy(aLocal, aGm[offset], this->tileLength * nLength);
        inQueueA.EnQue(aLocal);

        auto bLocal = inQueueB.AllocTensor<half>();
        DataCopy(bLocal, bGm[offset], this->tileLength * nLength);
        inQueueB.EnQue(bLocal);

        auto rLocal = inQueueR.AllocTensor<half>();
        DataCopy(rLocal, rGm[offset], this->tileLength * nLength);
        inQueueR.EnQue(rLocal);

    }

    __aicore__ inline void Compute(){
        LocalTensor<half> orih0Local = inQueueH.DeQue<half>();
        LocalTensor<half> oriwLocal = inQueueW.DeQue<half>();
        LocalTensor<half> orirLocal = inQueueR.DeQue<half>();
        LocalTensor<half> oriaLocal = inQueueA.DeQue<half>();
        LocalTensor<half> oribLocal = inQueueB.DeQue<half>();
        LocalTensor<half> orivLocal = inQueueV.DeQue<half>();
        LocalTensor<half> orikLocal = inQueueK.DeQue<half>();

        LocalTensor<half> vkLocal = vkBuf.Get<half>();
        LocalTensor<half> habLocal = habBuf.Get<half>();
        LocalTensor<float> tmpFp = tmpBuf.Get<float>();
        LocalTensor<float> calcFp = calcBuf.Get<float>();

        LocalTensor<half> orioLocal = outQueueO.AllocTensor<half>();
        LocalTensor<half> orihtLocal = outQueueH.AllocTensor<half>();

        LocalTensor<half> h0Local = orih0Local;
        LocalTensor<half> wLocal = oriwLocal;
        LocalTensor<half> rLocal = orirLocal;
        LocalTensor<half> aLocal = oriaLocal;
        LocalTensor<half> bLocal = oribLocal;
        LocalTensor<half> vLocal = orivLocal;
        LocalTensor<half> kLocal = orikLocal;

        LocalTensor<half> oLocal = orioLocal;
        LocalTensor<half> htLocal = orihtLocal;

        Exp(oriwLocal, oriwLocal, this->tileLength * nLength);
    
        #if __CCE_AICORE__ == 200
            uint64_t dstLocalList[4][16];
            for(int j = 0; j < 4; j++){
                for (int i = 0; i < 16; i++) {
                    dstLocalList[j][i] = (uint64_t)(calcFp[i * 32 + j * 8].GetPhyAddr());
                }
            }
            uint64_t srcLocalList[16];
            for (int i = 0; i < 16; i++) {
                srcLocalList[i] = (uint64_t)(tmpFp[0].GetPhyAddr());
            }
        #endif
        for(int i = 0; i < tileLength; i++){
            Mul(habLocal, h0Local, aLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            Cast(tmpFp, habLocal, RoundMode::CAST_NONE, nnLength);
            WholeReduceSum(tmpFp, tmpFp, nLength, nSize, 1, 1, 8); // get sa

            #if __CCE_AICORE__ == 200
                TransDataTo5HD<float>(dstLocalList[0], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[1], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[2], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[3], srcLocalList, {false, false, 8, 64, 1});
                Cast(htLocal, calcFp, RoundMode::CAST_NONE, nnLength);
                Mul(habLocal, htLocal, bLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });

                Cast(tmpFp, vLocal, RoundMode::CAST_NONE, nLength);
                TransDataTo5HD<float>(dstLocalList[0], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[1], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[2], srcLocalList, {false, false, 8, 64, 1});
                TransDataTo5HD<float>(dstLocalList[3], srcLocalList, {false, false, 8, 64, 1});
                Cast(htLocal, calcFp, RoundMode::CAST_NONE, nnLength);
            #else
                Brcb(calcFp, tmpFp, 8, {1, 8});
                Brcb(tmpFp, calcFp, 64, {1, 8});
                Cast(htLocal, tmpFp, RoundMode::CAST_NONE, nnLength);
                Mul(habLocal, htLocal, bLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });

                Cast(tmpFp, vLocal, RoundMode::CAST_NONE, nLength);
                Brcb(calcFp, tmpFp, 8, {1, 8});
                Brcb(tmpFp, calcFp, 64, {1, 8});
                Cast(htLocal, tmpFp, RoundMode::CAST_NONE, nnLength);
            #endif
            
            Mul(vkLocal, htLocal, kLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            
            Mul(h0Local, h0Local, wLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 }); // h0 = h0 * w
            Add(h0Local, h0Local, vkLocal, nnLength); // h0 = h0 + v*k
            Add(h0Local, h0Local, habLocal, nnLength); // h0 = h0 + hab

            // n,n n,1 -> n,1
            Mul(htLocal, h0Local, rLocal, nLength, nSize, { 1, 1, 1, 4, 4, 0 });
            Cast(tmpFp, htLocal, RoundMode::CAST_NONE, nnLength);
            WholeReduceSum(tmpFp, tmpFp, nLength, nSize, 1, 1, 8);
            Cast(oLocal, tmpFp, RoundMode::CAST_NONE, nLength);

            h0Local = h0Local[nnLength];
            wLocal = wLocal[nLength];
            rLocal = rLocal[nLength];
            aLocal = aLocal[nLength];
            bLocal = bLocal[nLength];
            vLocal = vLocal[nLength];
            kLocal = kLocal[nLength];
            oLocal = oLocal[nLength];
            htLocal = htLocal[nnLength];
        }
        
        inQueueA.FreeTensor(oriaLocal);
        inQueueB.FreeTensor(oribLocal);
        inQueueK.FreeTensor(orikLocal);
        inQueueV.FreeTensor(orivLocal);
        inQueueR.FreeTensor(orirLocal);
        inQueueW.FreeTensor(oriwLocal);
        outQueueH.FreeTensor(orihtLocal);

        inQueueH.EnQue<half>(orih0Local);
        outQueueO.EnQue<half>(orioLocal);
    }

    __aicore__ inline void CopyH0ToHt(){
        LocalTensor<half> h0Local = inQueueH.DeQue<half>();
        LocalTensor<half> htLocal = outQueueH.AllocTensor<half>();
        
        Adds(htLocal, h0Local, (half)0.0f, this->tileLength * nnLength);

        inQueueH.FreeTensor(h0Local);
        outQueueH.EnQue<half>(htLocal);
    }

    __aicore__ inline void CopyOutHt(int32_t offset)
    {
        // H, N, N(40, 64, 64)
        auto htLocal = outQueueH.DeQue<half>();
        DataCopy(htGm[offset], htLocal, this->tileLength * nnLength);
        outQueueH.FreeTensor(htLocal);
    }

    __aicore__ inline void CopyOut(int32_t offset)
    {
        auto oLocal = outQueueO.DeQue<half>();
        DataCopy(oGm[offset], oLocal, this->tileLength * nLength);
        outQueueO.FreeTensor(oLocal);
    }

private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueV, inQueueK, inQueueW, 
                inQueueR, inQueueA, inQueueB, inQueueH;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueO, outQueueH;
    GlobalTensor<half> kGm, vGm, wGm, rGm, aGm, bGm, oGm, h0Gm, htGm;
    TBuf<QuePosition::VECCALC> vkBuf, habBuf, tmpBuf, calcBuf;
    uint32_t tSize, bSize, thnLength, bhLength, tileLength, hSize, hnLength, hnnLength;
};

extern "C" __global__ __aicore__ void wkv7(GM_ADDR k, GM_ADDR v, GM_ADDR w, GM_ADDR r, 
                                            GM_ADDR a, GM_ADDR b, GM_ADDR h0, GM_ADDR o, GM_ADDR ht, 
                                            GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    #if __CCE_AICORE__ == 200
        KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_VECTOR_CORE);
    #endif
    TPipe pipe;
    if constexpr (TILING_KEY_IS(1)) {
        KernelWKV7TOne op;
        op.Init(tiling_data.bSize, tiling_data.tSize, tiling_data.hSize,
                k, v, w, r, a, b, h0, o, ht, &pipe);
        op.Process();
    }
    if constexpr (TILING_KEY_IS(2)){
        KernelWKV7 op;
        op.Init(tiling_data.bSize, tiling_data.tSize, tiling_data.hSize,
                k, v, w, r, a, b, h0, o, ht, &pipe);
        op.Process();
    }
}